#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2018 crane <crane@his-pc>
#
# Distributed under terms of the MIT license.

"""

"""


class TreeNode:
    def __init__(self, val):
        self.val = val
        self.left, self.right = None, None


from collections import defaultdict
from math import log


class Solution:
    def lowestCommonAncestor(self, root, node_a, node_b):
        # return self.regular_lowest_common_ancestor(root, node_a, node_b)
        # return self.tarjan_common_ancestor(root, node_a, node_b)
        # return self.rmq_common_ancestor(root, node_a, node_b)
        return self.doubling_tree(root, node_a, node_b)

    # =================== 树上倍增方法 =====================
    def doubling_tree(self, root, node_a, node_b):
        if not root:
            return None

        self.sparse_parent = defaultdict(dict)      # {node : {i:2^i'sparent}}
        self.depth = {}                             # {node: depth}
        self.node_a, self.node_b = node_a, node_b

        self.match = set([node_a, node_b])

        return self.dfs_double_tree(root, 0, None)

    def dfs_double_tree(self, node, depth, parent):
        self.depth[node] = depth
        if depth > 0:
            self._calc_all_sparse_parents(node, depth, parent)

        if node in self.match:
            self.match.remove(node)
            # self.match.add(node)

        if len(self.match) == 0:
            # print('match')

            # from pprint import pprint
            # pprint(self.sparse_parent)
            # pprint(self.depth)

            return self._same_parent(self.node_a, self.node_b)

        for child in [node.left, node.right]:
            if child:
                lca = self.dfs_double_tree(child, depth+1, node)
                if lca:
                    return lca

        return None

    def _calc_all_sparse_parents(self, node, depth, parent):
        self.sparse_parent[node][0] = parent    # 2 ** 0 == 1, 第一个parent
        pre_parent = parent
        base = int( log(depth, 2) )

        for i in range(1, base+1):
            # i_1_parent == pre_parent
            cur_parent = self.sparse_parent[pre_parent][i-1]
            self.sparse_parent[node][i] = cur_parent
            pre_parent = cur_parent

    def _same_parent(self, node_a, node_b):
        depth_a = self.depth[node_a]
        depth_b = self.depth[node_b]

        if  depth_a > depth_b:
            great_depth  = depth_a
            little_depth = depth_b
            great_node = node_a
            little_node = node_b
        else:
            great_depth  = depth_b
            little_depth = depth_a
            great_node = node_b
            little_node = node_a

        great_node = self._depth_parent(great_node, great_depth, little_depth)
        return self._same_depth_parent(great_node, little_node, little_depth)

    def _same_depth_parent(self, node_a, node_b, cur_depth):
        if node_b is node_a:
            return node_a

        print('same')
        parent_a, parent_b = node_a, node_b
        top_depth = 0
        bot_depth = cur_depth

        base = int(log(cur_depth, 2))

        while base >= 0:
            assert bot_depth >= top_depth

            parent_a = self.sparse_parent[parent_a][base]
            parent_b = self.sparse_parent[parent_b][base]

            move_depth = bot_depth - 2 ** base
            if parent_a is not parent_b:
                bot_depth = move_depth
                base = int(log(bot_depth-top_depth))
            else:
                top_depth = move_depth
                base -= 1

        # print(parent_a.val)
        print(parent_a)
        return parent_a

    def _depth_parent(self, node, cur_depth, target_depth):
        # node位于target_depth的parent
        diff = cur_depth - target_depth
        parent = node

        while diff > 0:
            i = int(log(diff, 2))
            parent = self.sparse_parent[parent][i]

            diff -= 2 ** i

        # print(target_depth)
        # print(parent.val)
        # print(node.val)

        return parent

    # =================== rmq method find lowest common ancestor =====================
    def rmq_common_ancestor(self, root, node_a, node_b):
        # 返回公共节点
        if not root:
            return None

        self.idx = 0
        self.occurence = []         # 通过索引查节点
        self.depth = []             # RMQ 最小depth的索引
        self.node_first = {}        # {node/node.value/id(node): idx}, 用户查询node第一次出现在occurence中的索引
                                    # 通过节点查索引

        self.dfs_build_depth_etc(root, depth=0)

        # build seg tree
        seg_size = len(self.depth) * 4
        self.seg_tree = seg_size * [None]
        self.build_segment_tree_idx(self.seg_tree, 0, self.depth, 0, len(self.depth))

        #
        idx_a = self.node_first[node_a]
        idx_b = self.node_first[node_b]
        little = min(idx_a, idx_b)
        great  = max(idx_a, idx_b)

        lca_idx = self.query_min_idx(little, great, self.depth, self.seg_tree)
        return self.occurence[lca_idx]

    def dfs_build_depth_etc(self, node, depth):
        # 中序遍历: NOTE: 但如果是一颗非二叉树, 这里则无法中序遍历.
        if node.left:
            self.dfs_build_depth_etc(node.left, depth+1)

        self.occurence.append(node)
        self.depth.append(depth)
        self.node_first[node] = self.idx
        self.idx += 1

        if node.right:
            self.dfs_build_depth_etc(node.right, depth+1)

    def build_segment_tree_idx(self, seg_tree, seg_idx, nums, start, end):
        # [start, end) 开闭区间
        assert end > start

        if start + 1 == end:
            seg_tree[seg_idx] = start
            return seg_tree[seg_idx]

        mid = (start + end + 1) // 2
        left_min_idx  = self.build_segment_tree_idx(seg_tree, seg_idx*2+1, nums, start, mid)
        right_min_idx = self.build_segment_tree_idx(seg_tree, seg_idx*2+2, nums, mid, end)

        if nums[left_min_idx] <= nums[right_min_idx]:
            seg_tree[seg_idx] = left_min_idx
        else:
            seg_tree[seg_idx] = right_min_idx

        return seg_tree[seg_idx]

    def query_min_idx(self, start, end, nums, seg_tree):
        # 询问[start, end]闭区间最小值在nums中索引
        nums_len = len(nums)
        assert end < nums_len

        # 开区间变闭区间
        return self._query_rec_idx(start, end+1, seg_tree, 0, 0, nums_len, nums)

    def _query_rec_idx(self, start, end, seg_tree, seg_idx, tree_start, tree_end, nums):
        if start <= tree_start and end >= tree_end:
            # query overlap seg tree: 完全相关
            return seg_tree[seg_idx]

        elif start >= tree_end or end <= tree_start:
            # query no overlap seg tree: 完全无关
            return None       # -1索引, 即最后一个元素中存放着一个很大的数
        else:
            # query partial overlap seg tree(可能seg tree rage overlap query rage)
            mid = (tree_start + tree_end + 1) // 2
            left_min_idx  = self._query_rec_idx(start, end, seg_tree, seg_idx*2+1, tree_start, mid, nums)
            right_min_idx = self._query_rec_idx(start, end, seg_tree, seg_idx*2+2, mid, tree_end, nums)
            # return min(left_min, right_min)
            if left_min_idx is None:
                return right_min_idx

            if right_min_idx is None:
                return left_min_idx

            if nums[left_min_idx] <= nums[right_min_idx]:
                return left_min_idx
            else:
                return right_min_idx

    # =================== tarjan算法 =====================
    def tarjan_common_ancestor(self, root, node_a, node_b):
        # 当前题目公共祖先只需要寻找一个pair, pair_list是为了处理多个pair的情况
        # NOTE: 并且tarjan公共祖先不一定是针对二叉树, 多叉树也是可以的
        va = node_a.val
        vb = node_b.val
        query_pairs = {va:vb, vb:va}
        # print(query_pairs)
        self.parent = {}
        # self.result = {}        # { (pair) : ancestor }
        self.result = None
        self.visited = set()
        self.query_pairs = query_pairs

        # do tarjan
        self.tarjan_dfs(root)

        # return list(self.result.values())[0]
        return self.result

    def union_find(self, val):
        # 递归find 和路径压缩(因为这里仅仅查找两个节点的公共祖先: 可以去掉路径压缩)
        parent_node = self.parent[val]
        parent_val = parent_node.val

        # 注意这里是if而不是while
        # if parent_val != val:
        #     self.parent[val] = self.union_find(parent_val)

        while parent_val != val:
            # 去掉了递归压缩路径
            val = parent_val
            parent_node = self.parent[val]
            parent_val  = parent_node.val

        return parent_node

    def tarjan_dfs(self, cur_node):
        val = cur_node.val
        self.parent[val] = cur_node
        self.visited.add(val)

        the_pair_val = self.query_pairs.pop(val, None)
        # if the_pair_val is not None:
        #     print('the_pair_val %s' % the_pair_val)
        # else:
        #     print('no pair of %s' % val)

        if the_pair_val in self.visited:
            # print('the_pair_val visited %s' % the_pair_val)
            # print(self.parent)
            # 这时候一定可以找到公共祖先(通过并查集)
            common_ancestor = self.union_find(the_pair_val)
            self.result = common_ancestor
            # print('query_pairs %s' % self.query_pairs)
            # print('the_pair not visited %s' % the_pair_val)

        if not self.query_pairs:
            return

        # if self.left:
        #     self.tarjan_dfs(cur_node.left)
        for child in [cur_node.left, cur_node.right]:
            if not child:
                continue

            self.tarjan_dfs(child)

            if not self.query_pairs:
                # query已经完毕, 不必再遍历后面的child了
                return

            self.parent[child.val] = cur_node

    # =================== 自己的方法 =====================
    def regular_lowest_common_ancestor(self, root, node_a, node_b):
        # assert va and vb in tree root
        if not root:
            return root

        self.paths = []

        ancestor = root
        va, vb = node_a.val, node_b.val
        self._find(root, [va, vb], [])
        print(self.paths)
        return self._find_common_path()

        # return self.tarjan_common_ancestor(root, {va:vb, vb:va})

    def _find(self, root, values, path):
        if not root:
            return None

        new_path = path + [root]
        for i in reversed(range(len(values))):
            v = values[i]
            if v == root.val:
                self.paths.append(new_path)
                values.pop(i)

        if not values:
            return

        self._find(root.left, values, new_path)
        self._find(root.right, values, new_path)

    def _find_common_path(self, ):
        min_len = min(map(len, self.paths))
        common = None
        for i in range(min_len):
            if not self.paths[0][i] is self.paths[1][i]:
                break
            common = self.paths[0][i]

        return common


# # for binary search tree
# while ancestor:
#     # if ancestor.val == va
#     v_ancestor = ancestor.val
#     if va == v_ancestor or vb == v_ancestor:
#         return ancestor
#     elif (va < v_ancestor and vb > v_ancestor) or (va > v_ancestor and va < v_ancestor):
#         return ancestor
#     elif va < v_ancestor and vb < v_ancestor:
#         ancestor = ancestor.left
#     else:
#         assert va > v_ancestor and vb > v_ancestor
#         ancestor = ancestor.right

# assert False, 'value must in Tree'
# # return ancestor


def main():
    print("start main")

    n1 = TreeNode(5)

    nl = TreeNode(0)
    nr = TreeNode(2)

    ll = TreeNode(3)

    n1.left, n1.right = nl, nr
    nl.left = ll

    s = Solution()
    # ret = s.lowestCommonAncestor(n1, ll, nr)
    ret = s.lowestCommonAncestor(ll, ll, ll)
    print(ret.val)

if __name__ == "__main__":
    main()
